{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "983yjns496tx"
   },
   "source": [
    "[![Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/pixeltable/pixeltable/blob/master/docs/release/tutorials/working-with-fireworks.ipynb)&nbsp;&nbsp;\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pixeltable/pixeltable/blob/master/docs/release/tutorials/working-with-fireworks.ipynb)\n",
    "\n",
    "# Working with Fireworks AI in Pixeltable\n",
    "\n",
    "Pixeltable's Fireworks integration enables you to access advanced large language models (LLMs) hosted on the Fireworks platform.\n",
    "\n",
    "### Prerequisites\n",
    "\n",
    "- A Fireworks account with an API key (https://fireworks.ai/api-keys)\n",
    "\n",
    "### Important Notes\n",
    "\n",
    "- Fireworks usage may incur costs based on your Fireworks plan.\n",
    "- Be mindful of sensitive data and consider security measures when integrating with external services.\n",
    "\n",
    "First you'll need to install required libraries and enter a Fireworks API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9pckrD01ik-e",
    "outputId": "060b8b32-48a6-48a0-e720-4eacf94d83ef"
   },
   "outputs": [],
   "source": [
    "%pip install -q pixeltable fireworks-ai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AQ6_Py7_7d0r",
    "outputId": "f82cfe36-be9e-4d43-f13e-9f6f5b680e8e"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import getpass\n",
    "\n",
    "if 'FIREWORKS_API_KEY' not in os.environ:\n",
    "    os.environ['FIREWORKS_API_KEY'] = getpass.getpass('Fireworks API Key:')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's create a Pixeltable directory to hold the tables for our demo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "brtjK-88tTSS",
    "outputId": "55d08c91-438a-4c3e-c217-3cea72faca11"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Connected to Pixeltable database at: postgresql://postgres:@/pixeltable?host=/Users/asiegel/.pixeltable/pgdata\n",
      "Created directory `fireworks_demo`.\n"
     ]
    }
   ],
   "source": [
    "import pixeltable as pxt\n",
    "\n",
    "pxt.drop_dir('fireworks_demo', force=True)  # Ensure a clean slate for the demo\n",
    "pxt.create_dir('fireworks_demo')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0kmjJoDq9Oqe"
   },
   "source": [
    "## Completions\n",
    "\n",
    "Create a Table: In Pixeltable, create a table with columns to  represent your input data and the columns where you want to store the results from Fireworks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5ti10tXu5m3X",
    "outputId": "30848066-1e9b-4efd-aad7-b2271a031ec3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created table `chat`.\n",
      "Added 0 column values with 0 errors.\n"
     ]
    }
   ],
   "source": [
    "from pixeltable.functions.fireworks import chat_completions\n",
    "\n",
    "# Create a table in Pixeltable and pick a model hosted on Fireworks with some parameters\n",
    "\n",
    "t = pxt.create_table('fireworks_demo.chat', {'input': pxt.StringType()})\n",
    "\n",
    "messages = [{'role': 'user', 'content': t.input}]\n",
    "t['output'] = chat_completions(\n",
    "    messages=messages,\n",
    "    model='accounts/fireworks/models/llama-v2-7b-chat',\n",
    "    # These parameters are optional and can be used to tune model behavior:\n",
    "    max_tokens=300,\n",
    "    top_k=40,\n",
    "    top_p=0.9,\n",
    "    temperature=0.7\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mePjoku95iUn",
    "outputId": "07c3f8f0-5301-44ba-ba33-25d10c3c26fa"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Added 0 column values with 0 errors.\n"
     ]
    }
   ],
   "source": [
    "# Parse the bot_response into a new column\n",
    "t['response'] = t.output.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 599
    },
    "id": "IkMM7OYb5rQ_",
    "outputId": "8e94af3e-485c-49f2-d7ba-b5490ec83af9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Computing cells: 100%|████████████████████████████████████████████| 3/3 [00:00<00:00,  3.35 cells/s]\n",
      "Inserting rows into `chat`: 1 rows [00:00, 236.35 rows/s]\n",
      "Computing cells: 100%|████████████████████████████████████████████| 3/3 [00:00<00:00,  3.32 cells/s]\n",
      "Inserted 1 row with 0 errors.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>input</th>\n",
       "      <th>response</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>Can you tell me who&#x27;s the President of the US?</td>\n",
       "      <td>Yes, I can certainly help you with that! The President of the United States is Joe Biden. He was inaugurated as the 46th President of the United States on January 20, 2021. Prior to his presidency, Biden served as the 47th Vice President of the United States under President Barack Obama from 2009 to 2017.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "                                            input  \\\n",
       "0  Can you tell me who's the President of the US?   \n",
       "\n",
       "                                            response  \n",
       "0    Yes, I can certainly help you with that! The...  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Start a conversation\n",
    "t.insert(input=\"Can you tell me who's the President of the US?\")\n",
    "t.select(t.input, t.response).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lTtQcjKQAlis"
   },
   "source": [
    "### Learn More\n",
    "\n",
    "To learn more about advanced techniques like RAG operations in Pixeltable, check out the [RAG Operations in Pixeltable](https://pixeltable.readme.io/docs/rag-operations-in-pixeltable) tutorial.\n",
    "\n",
    "If you have any questions, don't hesitate to reach out."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
